import os
from sklearn.externals import joblib
import sys
import numpy
from scipy.fftpack import dct
from scipy.fftpack import idct
import itertools 

if len(sys.argv) < 4 or len(sys.argv) > 5:
    print "Error: Wrong number of arguments."
    print "Please use python classifier.py #uniref #go_prior protein_list.txt protein_list_2.txt"
    print "uniref 0 - dont use uniref"
    print "uniref 1 - use uniref"
    print "go_prior 0 - proteins without go association are classified in default classifier"
    print "go_prior 1 - proteins without go association are classified as not interacting"
    sys.exit()

uniref = int(sys.argv[1])
go_prior = int(sys.argv[2])

if not (uniref == 0 or uniref ==1):    
    print "Error: Wrong number of arguments."
    print "Please use python classifier.py protein_list.txt protein_list_2.txt #uniref #go_prior"
    print "uniref 0 - dont use uniref"
    print "uniref 1 - use uniref"
    print "go_prior 0 - proteins without go association are classified in default classifier"
    print "go_prior 1 - proteins without go association are classified as not interacting"
    sys.exit()
if not (go_prior == 0 or go_prior ==1):    
    print "Error: Wrong number of arguments."
    print "Please use python classifier.py protein_list.txt protein_list_2.txt #uniref #go_prior"
    print "uniref 0 - dont use uniref"
    print "uniref 1 - use uniref"
    print "go_prior 0 - proteins without go association are classified in default classifier"
    print "go_prior 1 - proteins without go association are classified as not interacting"
    sys.exit()


f_go_ignore = open('ignore_gos','r');
GO_ignore =  f_go_ignore.readline().split()
f_go_ignore.close()

uniref_dic_protein_cluster = {}
uniref_dic_cluster_protein = {}

if (uniref == 1):

    f_uniref = open('uniref_aux.txt','r')
    line = f_uniref.readline()
    while (line !=''):
        try:
            aux = line.split()
            cluster = aux[1].split('_')[1]
            
            if cluster in uniref_dic_cluster_protein:
                uniref_dic_cluster_protein[cluster].append(aux[0])
            else:
                uniref_dic_cluster_protein[cluster] = []
                uniref_dic_cluster_protein[cluster].append(aux[0])
            
            if aux[0] in uniref_dic_protein_cluster:
                uniref_dic_protein_cluster[aux[0]].append(cluster)
            else:
                uniref_dic_protein_cluster[aux[0]] = []
                uniref_dic_protein_cluster[aux[0]].append(cluster)
        
            
            line = f_uniref.readline()
        except:
            error = 1;
    f_uniref.close()


genetic_code_dictionary = {}

def read_fasta(fp):
    name, seq = None, []
    for line in fp:
        line = line.rstrip()
        if line.startswith(">"):
            if name: yield (name, ''.join(seq))
            name, seq = line, []
        else:
            seq.append(line)
    if name: yield (name, ''.join(seq))
    
l1 = []
l2 = []
   
if len(sys.argv) == 4:
    genetic_code_dictionary1 = {}
    genetic_code_dictionary2 = {}
        
    with open(sys.argv[3]) as fp:
        for name, seq in read_fasta(fp):
            name = name.split("|")[1]
            if name not in genetic_code_dictionary1.keys():
                genetic_code_dictionary1[name] = seq
                
    l1 = sorted(genetic_code_dictionary1.keys())
    
    genetic_code_dictionary = dict(genetic_code_dictionary1.items())
    
if len(sys.argv) == 5:
    #duas listas            
    genetic_code_dictionary1 = {}
    genetic_code_dictionary2 = {}
        
    with open(sys.argv[3]) as fp:
        for name, seq in read_fasta(fp):
            name = name.split("|")[1]
            if name not in genetic_code_dictionary1.keys():
                genetic_code_dictionary1[name] = seq
                
    l1 = sorted(genetic_code_dictionary1.keys())
            
    with open(sys.argv[4]) as fp:
        for name, seq in read_fasta(fp):
            name = name.split("|")[1]
            if name not in genetic_code_dictionary2.keys():
                genetic_code_dictionary2[name] = seq
                
    l2 = sorted(genetic_code_dictionary2.keys())
    
    genetic_code_dictionary = dict(genetic_code_dictionary1.items() + genetic_code_dictionary2.items())
    


# -- Load GOs --
GO_dict = {}
GOs = ['go%3a1071','go%3a3824','go%3a4872','go%3a5085','go%3a5198','go%3a5215','go%3a5488','go%3a9055','go%3a16015','go%3a16209','go%3a16247','go%3a16530','go%3a30234','go%3a30545','go%3a31386','go%3a36370','go%3a42056','go%3a45182','go%3a45499','go%3a45735','go%3a60089','go%3a5515']
for go in GOs:
    f_go = open("uniprot-"+go+".list", "r")
    prot_l = f_go.readlines()
    
    for elem in prot_l:
        prot = elem.rstrip()
        if prot in GO_dict:
            GO_dict[prot].append(go)
        else:
            GO_dict[prot] = []
            GO_dict[prot].append(go)
    print go    
    f_go.close()


GO_classifiers = {}
for f in os.listdir(os.getcwd()):
    if f.endswith("_classifier"):
        clf2 = joblib.load(f)
        GO_classifiers[f] = clf2
        print  GO_classifiers[f]
GO_Datasets = {}

def substitution(amino_acids):
    newChain= []
    
    for amino_acid in amino_acids:
        if amino_acid == 'A' or amino_acid == 'G' or amino_acid == 'V' :
            newChain.append("0");
        elif amino_acid == 'I' or amino_acid == 'L' or amino_acid == 'F' or amino_acid == 'P' :
            newChain.append("1");
        elif amino_acid == 'Y' or amino_acid == 'M' or amino_acid == 'T' or amino_acid == 'S' :
            newChain.append("2");
        elif amino_acid == 'H' or amino_acid == 'N' or amino_acid == 'Q' or amino_acid == 'W' :
            newChain.append("3");
        elif amino_acid == 'R' or amino_acid == 'K':
            newChain.append("4");
        elif amino_acid == 'D' or amino_acid == 'E':
            newChain.append("5");
        elif amino_acid == 'C':
            newChain.append("6");
    
    chain = ','.join(newChain);
    return chain;

def getPCChain(amino_acids):
    chain = substitution(amino_acids)
    return chain

protein_features = {}

inx = 0

if len(sys.argv) == 4:
    
    for e in itertools.combinations(l1,2):
        
        protein_interaction =  (e[0].strip()+" "+e[1].strip()).split()
        
        protein_1 = protein_interaction[0];
        protein_2 = protein_interaction[1];
        
        contains = protein_features.keys()
    
        if protein_1 not in contains:
    
            amino_acids_protein_1 = genetic_code_dictionary[protein_1];
            chain_protein_1 = getPCChain(amino_acids_protein_1);
            
            prot_features = chain_protein_1.split(",")
            valores = map(float, prot_features)
            
            valores = numpy.array(valores,dtype='float')
            
            J = dct(valores)
            features = numpy.zeros(600)
            J= J[:600]
            
            if len(features) < len(J):
                c = J.copy()
                c[:len(features)] += features
            else:
                c = features.copy()
                c[:len(J)] += J
            reconstructed = idct(c)
            
            media = numpy.average(reconstructed);
            desvioPadrao = numpy.std(reconstructed)
            for j in range(0,len(reconstructed)):
                reconstructed[j] = (reconstructed[j]-media)/desvioPadrao;
                
            protein_features[protein_1] = reconstructed; 
            contains.append(protein_1)
        
        if protein_2 not in contains:
            amino_acids_protein_2 = genetic_code_dictionary[protein_2];
            chain_protein_2 = getPCChain(amino_acids_protein_2);
            
            prot_features = chain_protein_2.split(",")
        
            valores = map(float, prot_features)
            
            valores = numpy.array(valores,dtype='float')
    
            J = dct(valores)
            features = numpy.zeros(600)
            J= J[:600]
        
            if len(features) < len(J):
                c = J.copy()
                c[:len(features)] += features
            else:
                c = features.copy()
                c[:len(J)] += J
            reconstructed = idct(c)
            
            media = numpy.average(reconstructed);
            desvioPadrao = numpy.std(reconstructed)
            for j in range(0,len(reconstructed)):
                reconstructed[j] = (reconstructed[j]-media)/desvioPadrao;
                
                
            protein_features[protein_2] = reconstructed; 
            contains.append(protein_2)
        
        encontra = 0
        
        if protein_1 in GO_dict and protein_2 in GO_dict:
            probas = []
            encontra = 0
            for k in GO_dict[protein_1]:
                for j in GO_dict[protein_2]:
                    key = k+j+"_classifier";
                    key2 = j+k+"_classifier";
            
                    if (key < key2):
                        key = key        
                    else:
                        key = key2
                    
                    if (go_prior == 1): 
                        if key in GO_ignore:
                            encontra = 1
                            probas.append(0)
                            break;
                        
                    if key in GO_classifiers:
                        encontra = 1
                        proba = GO_classifiers[key].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                        probas.append(proba) 
                                 
            if encontra == 1:
                print protein_1 +" "+protein_2+ " " +str(numpy.average(probas))
            else:
                if go_prior == 1:
                    pass
                else:
                    proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                    print protein_1 +" "+protein_2+ " " +str(proba)
            
        elif uniref == 1:
            found = 0
    
            try:
                cluster = uniref_dic_protein_cluster[protein_1]
                cluster2 = uniref_dic_protein_cluster[protein_2] 
                for p1 in uniref_dic_protein_cluster[cluster]:
                    for p2 in uniref_dic_protein_cluster[cluster2]:
                        if p1 in GO_dict and p2 in GO_dict:
                            found = 1
                            break;
            except:
                found = 0
                
            if found == 1:
                probas = []
                for k in GO_dict[protein_1]:
                    for j in GO_dict[protein_2]:
                        key = k+j+"_classifier";
                        key2 = j+k+"_classifier";
                
                        if (key < key2):
                            key = key        
                        else:
                            key = key2
                            
                        if (go_prior == 1): 
                            if key in GO_ignore:
                                encontra = 1
                                probas.append(0)
                                break;
                    
                        if key in GO_classifiers:
                            encontra = 1
                            proba = GO_classifiers[key].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                            probas.append(proba) 
                                     
                if encontra == 1:
                    print protein_1 +" "+protein_2+ " " +str(numpy.average(probas))
                else:
                    if go_prior == 1:
                        pass
                    else:
                        proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                        print protein_1 +" "+protein_2+ " " +str(proba)
    
            else:        
                if go_prior == 1:
                    pass
                else:
                    proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                    print protein_1 +" "+protein_2+ " " +str(proba)
        else:        
            if go_prior == 1:
                pass
            else:
                proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                print protein_1 +" "+protein_2+ " " +str(proba)
    
        if (inx %1000 == 0):
            protein_features = {}
        
        inx +=1
        
else:
    for i in l1:
        for j in l2:
            if i!=j:    
                protein_interaction = (i.strip()+" "+j.strip()).split()
                protein_1 = protein_interaction[0];
                protein_2 = protein_interaction[1];

                contains = protein_features.keys()
            
                if protein_1 not in contains:
            
                    amino_acids_protein_1 = genetic_code_dictionary[protein_1];
                    chain_protein_1 = getPCChain(amino_acids_protein_1);
                    
                    prot_features = chain_protein_1.split(",")
                    valores = map(float, prot_features)
                    
                    valores = numpy.array(valores,dtype='float')
                    
                    J = dct(valores)
                    features = numpy.zeros(600)
                    J= J[:600]
                    
                    if len(features) < len(J):
                        c = J.copy()
                        c[:len(features)] += features
                    else:
                        c = features.copy()
                        c[:len(J)] += J
                    reconstructed = idct(c)
                    
                    media = numpy.average(reconstructed);
                    desvioPadrao = numpy.std(reconstructed)
                    for j in range(0,len(reconstructed)):
                        reconstructed[j] = (reconstructed[j]-media)/desvioPadrao;
                        
                    protein_features[protein_1] = reconstructed; 
                    contains.append(protein_1)
                
                if protein_2 not in contains:
                    amino_acids_protein_2 = genetic_code_dictionary[protein_2];
                    chain_protein_2 = getPCChain(amino_acids_protein_2);
                    
                    prot_features = chain_protein_2.split(",")
                
                    valores = map(float, prot_features)
                    
                    valores = numpy.array(valores,dtype='float')
            
                    J = dct(valores)
                    features = numpy.zeros(600)
                    J= J[:600]
                
                    if len(features) < len(J):
                        c = J.copy()
                        c[:len(features)] += features
                    else:
                        c = features.copy()
                        c[:len(J)] += J
                    reconstructed = idct(c)
                    
                    media = numpy.average(reconstructed);
                    desvioPadrao = numpy.std(reconstructed)
                    for j in range(0,len(reconstructed)):
                        reconstructed[j] = (reconstructed[j]-media)/desvioPadrao;
                        
                        
                    protein_features[protein_2] = reconstructed; 
                    contains.append(protein_2)
                
                encontra = 0
                
                if protein_1 in GO_dict and protein_2 in GO_dict:
                    probas = []
                    encontra = 0
                    for k in GO_dict[protein_1]:
                        for j in GO_dict[protein_2]:
                            key = k+j+"_classifier";
                            key2 = j+k+"_classifier";
                    
                            if (key < key2):
                                key = key        
                            else:
                                key = key2
                            
                            if (go_prior == 1): 
                                if key in GO_ignore:
                                    encontra = 1
                                    probas.append(0)
                                    break;
                                
                            if key in GO_classifiers:
                                encontra = 1
                                proba = GO_classifiers[key].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                                probas.append(proba) 
                                         
                    if encontra == 1:
                        print protein_1 +" "+protein_2+ " " +str(numpy.average(probas))
                    else:
                        if go_prior == 1:
                            pass
                        else:
                            proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                            print protein_1 +" "+protein_2+ " " +str(proba)
                    
                elif uniref == 1:
                    found = 0
            
                    try:
                        cluster = uniref_dic_protein_cluster[protein_1]
                        cluster2 = uniref_dic_protein_cluster[protein_2] 
                        for p1 in uniref_dic_protein_cluster[cluster]:
                            for p2 in uniref_dic_protein_cluster[cluster2]:
                                if p1 in GO_dict and p2 in GO_dict:
                                    found = 1
                                    break;
                    except:
                        found = 0
                        
                    if found == 1:
                        probas = []
                        for k in GO_dict[protein_1]:
                            for j in GO_dict[protein_2]:
                                key = k+j+"_classifier";
                                key2 = j+k+"_classifier";
                        
                                if (key < key2):
                                    key = key        
                                else:
                                    key = key2
                                    
                                if (go_prior == 1): 
                                    if key in GO_ignore:
                                        encontra = 1
                                        probas.append(0)
                                        break;
                            
                                if key in GO_classifiers:
                                    encontra = 1
                                    proba = GO_classifiers[key].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                                    probas.append(proba) 
                                             
                        if encontra == 1:
                            print protein_1 +" "+protein_2+ " " +str(numpy.average(probas))
                        else:
                            if go_prior == 1:
                                pass
                            else:
                                proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                                print protein_1 +" "+protein_2+ " " +str(proba)
            
                    else:        
                        if go_prior == 1:
                            pass
                        else:
                            proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                            print protein_1 +" "+protein_2+ " " +str(proba)
                else:        
                    if go_prior == 1:
                        pass
                    else:
                        proba = GO_classifiers["no_go_classifier"].predict_proba(numpy.concatenate((protein_features[protein_1],protein_features[protein_2])))[0][1]
                        print protein_1 +" "+protein_2+ " " +str(proba)
            
                if (inx %1000 == 0):
                    protein_features = {}
                
                inx +=1